import java.io.File;
import java.io.IOException;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import edu.cmu.minorthird.text.MonotonicTextLabels;
import edu.cmu.minorthird.text.MutableTextLabels;
import edu.cmu.minorthird.text.Span;
import edu.cmu.minorthird.text.TextBaseLoader;
import edu.cmu.minorthird.text.TextLabels;
import edu.cmu.minorthird.text.learn.SpanFeatureExtractor;
import edu.umass.cs.mallet.base.fst.CRF2;
import edu.umass.cs.mallet.base.fst.TokenAccuracyEvaluator;
import edu.umass.cs.mallet.base.types.FeatureVectorSequence;
import edu.umass.cs.mallet.base.types.Instance;
import edu.umass.cs.mallet.base.types.InstanceList;
import edu.umass.cs.mallet.base.types.Sequence;
import etxt2db.features.CharacterFeatureClassifier;
import etxt2db.features.CharacterTypeFeatureClassifier;
import etxt2db.features.EditableTokenFE;
import etxt2db.features.PatternFeatureClassifier;
import etxt2db.features.ValueCaseInsensitiveFeatureClassifier;
import etxt2db.features.ValueCaseSensitiveFeatureClassifier;
import etxt2db.mallet.MalletLoader;


public class SolutionMalletCRF {
	private static String trainingSet = "C:\\Users\\Goncalo\\Desktop\\SeminarPartial\\";
	private static CRF2 crf;
	
	private static void convertSequenceTextLabels(MonotonicTextLabels text, Map<String,Sequence> labels){
		Iterator<Span> iter = text.getTextBase().documentSpanIterator();
		while(iter.hasNext()){
			Span current = iter.next();
			int i = 0;
			List<Span> list = new ArrayList<Span>();
			while(i<current.size()){
				Span tok = current.subSpan(i, 1);
				if(!tok.asString().matches("\\s*")){
					list.add(tok);
				}
				i++;
			}
			
			Sequence textLabels = labels.get(current.getDocumentId());
			
			int currentSpanBegin=0;
			int currentSpanEnd=0;
			
			
			for(int j=0; j<textLabels.size(); j++){
				if(((String)textLabels.get(j)).endsWith("Unique")){
					currentSpanBegin = list.get(j).getLoTextToken();
					currentSpanEnd = list.get(j).getLoTextToken()+1;
					Span newSpan=current.charIndexProperSubSpan(currentSpanBegin,currentSpanEnd);
					//text.addToType(newSpan,tag.substring(0, tag.lastIndexOf("Unique")));
					System.out.println(newSpan.asString());
				}else if(((String)textLabels.get(j)).endsWith("Begin")){
					currentSpanBegin = list.get(j).getLoTextToken();
				}else if(((String)textLabels.get(j)).endsWith("End")){
					currentSpanEnd = list.get(j).getLoTextToken()+1;
					Span newSpan=current.subSpan(currentSpanBegin, currentSpanEnd-currentSpanBegin);
					//text.addToType(newSpan,tag.substring(0, tag.lastIndexOf("End")));
					System.out.println(newSpan.asString());
				}
			}
		}
	}
	
	public static void main(String[] args) throws IOException, ParseException{
		List<CharacterFeatureClassifier> listFE = new ArrayList<CharacterFeatureClassifier>();
		listFE.add(new CharacterTypeFeatureClassifier());
		listFE.add(new PatternFeatureClassifier());
		listFE.add(new ValueCaseSensitiveFeatureClassifier());
		listFE.add(new ValueCaseInsensitiveFeatureClassifier());
		SpanFeatureExtractor featureExtractor = new EditableTokenFE(listFE);
		
		File treino = new File(trainingSet);
		TextBaseLoader loader = new TextBaseLoader();
		loader.load(treino);
		TextLabels treinoLabels = loader.getLabels();
		
		List<String> attributes = new ArrayList<String>();
		attributes.add("location");
		
		MalletLoader ml = new MalletLoader();
		InstanceList trainingData = ml.load(treinoLabels, attributes,featureExtractor);
		
		TokenAccuracyEvaluator eval = new TokenAccuracyEvaluator();
		
		crf = new CRF2(trainingData.getDataAlphabet(), trainingData.getTargetAlphabet());
		crf.addStatesForHalfLabelsConnectedAsIn (trainingData);
		crf.setGaussianPriorVariance (Double.POSITIVE_INFINITY);
		crf.train (trainingData, null, null, eval, 100);
		
		
		File teste = new File("./resources/test/cmu.andrew.official.cmu-news-2457-0");
		TextBaseLoader loaderTest = new TextBaseLoader();
		loaderTest.load(teste);
		MutableTextLabels testeLabels = loaderTest.getLabels();
		InstanceList testingData = ml.load(testeLabels, attributes,featureExtractor);
		
		Map<String,Sequence> result = new HashMap<String,Sequence>();
	    // Iterate over the sequences in the InstanceList.
	    for (InstanceList.Iterator i = testingData.iterator();i.hasNext();) {
	    	Instance ins = i.nextInstance();
	        FeatureVectorSequence sequence = (FeatureVectorSequence)ins.getData();
	        Sequence labels = crf.viterbiPath(sequence).output();
	        result.put((String)ins.getSource(), labels);
	    }
	    convertSequenceTextLabels(testeLabels, result);
	}

	
	
}
